Add ZImage LoRA support and integrate into ZImagePipeline #12750
Conversation
|
Tested using a simple script. Testing script#!/usr/bin/env python
"""Test script for ZImage LoRA support."""
import sys
sys.path.insert(0, '/home/ohiom/diffusers/src')
import torch
from diffusers import ZImagePipeline
# Paths
MODEL_PATH = "database/models/huggingface/models--Tongyi-MAI--Z-Image-Turbo/snapshots/78771b7e11b922c868dd766476bda1f4fc6bfc96"
LORA_PATH = "TechnicallyColorZ_V1.safetensors"
print("Loading ZImagePipeline...")
pipe = ZImagePipeline.from_pretrained(
MODEL_PATH,
torch_dtype=torch.bfloat16,
local_files_only=True,
)
pipe.to("cuda")
print(f"Pipeline loaded. Has load_lora_weights: {hasattr(pipe, 'load_lora_weights')}")
print(f"\nLoading LoRA from {LORA_PATH}...")
pipe.load_lora_weights(LORA_PATH)
print("LoRA loaded successfully!")
# Generate an image
prompt = "t3chnic4lly vibrant 1960s close-up of a woman sitting under a tree in a blue skit and white blouse, she has blonde wavy short hair and a smile with green eyes lake scene by a garden with flowers in the foreground 1960s styl;e film She's holding her hand out there is a small smooth frog in her palm, she's making eye contact with the toad."
print(f"\nGenerating image with prompt: {prompt}")
image = pipe(
prompt=prompt,
num_inference_steps=8,
guidance_scale=1.0,
height=1024,
width=1024,
generator=torch.Generator(device="cuda").manual_seed(42),
).images[0]
output_path = "test_zimage_lora_output.png"
image.save(output_path)
print(f"\nImage saved to {output_path}") |
|
@asomoza cc |
|
thanks for the contribution, it works ok without scale but we need to add the transformer model here With that:
Also I did a quick run on the tests to check if this was detected but they all fail with: |
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
What a clean PR this is! Thanks so much for contributing!
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
- Override test_lora_fuse_nan to use ZImage's 'layers' attribute instead of 'transformer_blocks' - Skip block-level LoRA scaling test (not supported in ZImage) - Add required imports: numpy, torch_device, check_if_lora_correctly_set
|
@sayakpaul You are too kind. As for the remaining failures, as far as I can tell, due to ZImage requiring |
|
thanks a lot @CalamitousFelicitousness, can you please add the ZImageLoraLoaderMixin to the lora docs here for the tests, I leave it to @sayakpaul to make the decision |
|
Hmm, the failing tests -- are these because of the numerical instabilities arising from the use of |
ZImage uses 'attention.to_k' naming convention instead of 'attn.to_k', so the base test's module name search loop never finds a match. This override uses the correct naming pattern for ZImage architecture.
|
Squashed one more naming convention issue with an override. My best somewhat-educated 2AM guess is that it lies somewhere between the Complex64 RoPE and torch.empty, unfortunately here I am beginning to approach the functional limits of my practical knowledge, so if anything comes to mind later on, I will bring it up. |
|
@CalamitousFelicitousness that is understandable. I am facing something similar in #12741 |
|
Let's try to add a |
|
They are really flaky. I'm doing 10 retries right now, and it whittled down the failures to 6 at the lowest, some tests are failing one run and pass on another. I will just set it to 100 and see if they pass. 100 retriesAnother 100 retriesStatus: We have three remaining tests that haven't passed yet:
|
|
@sayakpaul The last three tests fail, even at 250 retries, I tried to initialise the padding tokens to try and get them to pass, but no dice. |
4 participants
What does this PR do?
Adds support for LoRA for ZImagePipeline and adds a conversion script into Diffusers format. Resolves issue #12745
Before submitting
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@sayakpaul